#!/usr/bin/env ts-node

import { PrismaClient } from '@prisma/client';

async function updateSectorsForExistingStocks() {
  console.log('🔄 开始更新现有股票的板块分类信息...');
  
  const prisma = new PrismaClient();
  
  try {
    // 获取所有未分类的股票
    const unclassifiedStocks = await prisma.stock.findMany({
      where: {
        OR: [
          { sector: null },
          { sector: '' }
        ]
      },
      include: {
        market: true
      }
    });
    
    console.log(`📋 找到 ${unclassifiedStocks.length} 只未分类股票`);
    
    if (unclassifiedStocks.length === 0) {
      console.log('✅ 所有股票都已分类');
      return;
    }
    
    let updatedCount = 0;
    
    // 批量更新股票板块信息
    for (const stock of unclassifiedStocks) {
      try {
        let sector = '其他';
        
        // 根据股票代码确定板块分类
        if (stock.market.marketCode === 'SH') {
          // 沪市
          if (stock.stockCode.startsWith('688')) {
            sector = '科创板';
          } else if (stock.stockCode.startsWith('600') || 
                     stock.stockCode.startsWith('601') || 
                     stock.stockCode.startsWith('603') || 
                     stock.stockCode.startsWith('605')) {
            sector = '沪A主板';
          } else {
            sector = '沪市其他';
          }
        } else if (stock.market.marketCode === 'SZ') {
          // 深市
          if (stock.stockCode.startsWith('300')) {
            sector = '创业板';
          } else if (stock.stockCode.startsWith('000') || 
                     stock.stockCode.startsWith('001')) {
            sector = '深A主板';
          } else if (stock.stockCode.startsWith('002')) {
            sector = '中小板';
          } else {
            sector = '深市其他';
          }
        } else if (stock.market.marketCode === 'HK') {
          sector = '港股';
        } else if (stock.market.marketCode === 'US') {
          sector = '美股';
        }
        
        // 更新股票板块信息
        await prisma.stock.update({
          where: { id: stock.id },
          data: { 
            sector,
            updatedAt: new Date()
          }
        });
        
        updatedCount++;
        
        if (updatedCount % 50 === 0) {
          console.log(`📦 已更新 ${updatedCount}/${unclassifiedStocks.length} 只股票...`);
        }
        
      } catch (error: any) {
        console.error(`❌ 更新股票 ${stock.stockCode} 失败:`, error.message);
      }
    }
    
    console.log(`✅ 板块分类更新完成! 共更新 ${updatedCount} 只股票`);
    
    // 显示更新后的统计
    console.log('\n📊 更新后的板块分布:');
    const sectorStats = await prisma.stock.groupBy({
      by: ['sector'],
      _count: { _all: true }
    });
    
    for (const stat of sectorStats) {
      const sector = stat.sector || '未分类';
      console.log(`${sector}: ${stat._count._all} 只`);
    }
    
  } catch (error) {
    console.error('❌ 更新失败:', error);
  } finally {
    await prisma.$disconnect();
  }
}

updateSectorsForExistingStocks().then(() => {
  console.log('🎉 程序执行完成');
  process.exit(0);
}).catch(console.error);